#!/usr/bin/env bash
# Copyright (c) 2024 maminjie <canpool@163.com>
# SPDX-License-Identifier: MulanPSL-2.0

ROOT_DIR=$(dirname $(readlink -f "$0"))
CUR_DIR=$(pwd)
WORK_DIR=$CUR_DIR/iso_custom

MOUNT_DIR=$WORK_DIR/mount
BUILD_DIR=$WORK_DIR/build
ISO_DIR=$WORK_DIR/iso
IMG_DIR=$WORK_DIR/image

PRODUCT="MyOS"
VERSION="1.0"
RELEASE=""
ARCH="$(uname -m)"
ISO_INFILE=""
ISO_OUTFILE=""
OS_FULLNAME=""
LOGO_DIR=""

usage() {
printf "customize operating system from openeuler iso

usage:
    bash iso_custom.sh [-h] [-p PRODUCT] [-v VERSION] [-r RELEASE] [-i OLDISO] [-o NEWISO] [-g LOGO]

optional arguments:
    -p PRODUCT     product name, default is MyOS
    -v VERSION     version identifier, default is 1.0
    -r RELEASE     release information, default is null
    -i OLDISO      openeuler iso filepath
    -o NEWISO      new iso filename, default is PRODUCT-VERSION-RELEASE-ARCH-dvd.iso
    -g LOGO        directory of logos, directory structure is same as openEuler-logos.rpm
    -h             show the help message and exit
\n"
}


log_error() {
    echo "ERROR: $1"
}

log_info() {
    echo "INFO: $1"
}


parse_cmdline() {
    if [ $# -eq 0 ]; then
        usage; exit 0
    fi

    if [[ "${ARCH}" != "aarch64" && "${ARCH}" != "x86_64" ]]; then
        log_error "unsupported architecture: ${ARCH}"; return 1
    fi

    # parse cmd line arguments
    while getopts ":p:v:r:i:o:g:h" opt; do
        case "$opt" in
            p)
                PRODUCT="$OPTARG"
            ;;
            v)
                VERSION="$OPTARG"
            ;;
            r)
                RELEASE="$OPTARG"
            ;;
            i)
                ISO_INFILE="$OPTARG"
            ;;
            o)
                ISO_OUTFILE="$OPTARG"
            ;;
            g)
                LOGO_DIR="$(realpath $OPTARG)"
            ;;
            h)
                usage; exit 0
            ;;
            ?)
                echo "please check the params."; usage; return 1
            ;;
        esac
    done
    if [ ! -f "$ISO_INFILE" ]; then
        echo "invalid ISO arguments"; return 1
    fi
    if [ -z "$ISO_OUTFILE" ]; then
        if [ -z "$RELEASE" ]; then
            OS_FULLNAME="$PRODUCT-$VERSION"
        else
            OS_FULLNAME="$PRODUCT-$VERSION-$RELEASE"
        fi
        ISO_OUTFILE="$OS_FULLNAME-$ARCH-dvd.iso"
    fi
    return 0
}

check_user() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root!"
        exit 1
    fi
}

check_requires() {
    local ret=0

    # createrepo_c package contains:
    #   createrepo
    # rpmdevtools package contains:
    #   rpmdev-setuptree
    # rpmrebuild package contains:
    #   rpmrebuild
    # rpm-build package contains:
    #   rpmbuild
    # rpm package contains:
    #   rpm2cpio
    # squashfs-tools package contains:
    #   unsquashfs mksquashfs
    # kpartx package contains:
    #   kpartx
    # genisoimage package contains:
    #   mkisofs
    # isomd5sum package contains
    #   implantisomd5
    # eg:
    #   dnf install createrepo_c rpmdevtools rpmrebuild rpm-build genisoimage isomd5sum -y
    for c in rpmdev-setuptree rpm2cpio rpmrebuild rpmbuild createrepo unsquashfs kpartx mksquashfs mkisofs implantisomd5; do
        command -v $c > /dev/null
        if [ $? -ne 0 ]; then
            log_error "command $c not found"
            ret=1
        fi
    done

    return $ret
}

# make_iso iso_label iso_file iso_dir
# params:
#   iso_label - equal to LABEL in EFI/BOOT/grub.cfg
#   iso_file - new iso file name
#   iso_dir - direcotry of iso files
make_iso() {
    if [ $# -lt 3 ]; then
        echo "invalid param"; return 1
    fi

    local arch=$(uname -m)
    local iso_label=$1
    local iso_file=$2
    local iso_dir=$(realpath $3)

    if [ ! -d "$iso_dir" ]; then
        echo "$iso_dir: no such directory"; return 1
    fi

    if [[ "$arch" = "x86_64" ]]; then
        mkisofs -R -J -T -r -l -d -joliet-long -allow-multidot -allow-leading-dots -no-bak -V "${iso_label}" -o "${iso_file}" -b isolinux/isolinux.bin \
            -c isolinux/boot.cat -no-emul-boot -boot-load-size 4 -boot-info-table  -eltorito-alt-boot -e images/efiboot.img -no-emul-boot "${iso_dir}"
        [[ $? -ne 0 ]] && return 1
        isohybrid -u "$iso_file"
    elif [[ "$arch" = "aarch64" || "$arch" = "loongarch64" ]]; then
        mkisofs -R -J -T -r -l -d -joliet-long -allow-multidot -allow-leading-dots -no-bak -V "${iso_label}" -o "${iso_file}" -e images/efiboot.img -no-emul-boot "${iso_dir}"
        [[ $? -ne 0 ]] && return 1
    else
        echo "unsupported architecture: ${arch}"; return 1
    fi
    implantisomd5 "$iso_file"
    return 0
}

init_dir() {
    [[ ! -d "$1" ]] && mkdir -p "$1"
}

clean_dir() {
    [[ -d "$1" ]] && rm -rf "$1"
}

custom_iso() {
    local version_str=$VERSION
    local os_version=$VERSION
    if [ -n "$RELEASE" ]; then
        version_str="$VERSION ($RELEASE)"
        os_version="$VERSION-$RELEASE"
    fi
    local product=$(echo "$PRODUCT" | tr A-Z a-z)
    local oe_version="${ISO_INFILE#*-}"
    oe_version="${oe_version%%-*}"

    # init dirs
    init_dir $WORK_DIR
    init_dir $MOUNT_DIR
    init_dir $BUILD_DIR

    # mount iso
    umount $MOUNT_DIR 2> /dev/null
    mount $ISO_INFILE $MOUNT_DIR > /dev/null

    local rpm_dst_dir="/root/rpmbuild/RPMS"
    [[ ! -d "$rpm_dst_dir" ]] && rpmdev-setuptree
    rm -rf $rpm_dst_dir/*

    local rpm_src_dir="$MOUNT_DIR/Packages"
    local pkgs=$(ls $rpm_src_dir/*.rpm)

    # rebuild rpm packages
    for pkg in $pkgs; do
        local pkg_file_name=${pkg##*/}
        local pkg_name=$(echo "$pkg_file_name" | sed -r 's/-[^-]+-[^-]+$//')

        local target_pkgs=(
            openEuler-release openEuler-logos openEuler-latest-release
            anaconda-core python3-productmd
        )
        local is_target=0
        for t_pkg in ${target_pkgs[@]}; do
            if [ $pkg_name = $t_pkg ]; then
                is_target=1
                break
            fi
        done
        if [ $is_target -eq 0 ]; then
            continue
        fi

        local pkg_file=$pkg
        local pkg_build_dir=$BUILD_DIR/$pkg_name
        local spec_file=$pkg_build_dir/$pkg_name.spec
        # make build dir
        [[ -d $pkg_build_dir ]] && rm -rf $pkg_build_dir
        mkdir -p $pkg_build_dir
        # uncompress
        rpm2cpio $pkg_file | cpio -id -D $pkg_build_dir
        # extract spec
        rpmrebuild -n -p -s $spec_file $pkg_file
        # TODO: change spec vendor infos

        case "$pkg_name" in
            "openEuler-release")
                oe_version=$(cat $pkg_build_dir/etc/os-release | grep "VERSION=" | awk -F = '{print $2}' | sed -r 's/"|\)//g' | sed -r 's/\s+\(/-/g')
                cat << EOF > $pkg_build_dir/etc/os-release
NAME="$PRODUCT"
VERSION="$version_str"
ID="$PRODUCT"
VERSION_ID="$VERSION"
PRETTY_NAME="$PRODUCT $version_str"
ANSI_COLOR="0;31"
EOF
                mv $pkg_build_dir/etc/openEuler-release $pkg_build_dir/etc/$PRODUCT-release
                mv $pkg_build_dir/usr/share/doc/openEuler-release $pkg_build_dir/usr/share/doc/$PRODUCT-release
                mv $pkg_build_dir/usr/share/openEuler-release $pkg_build_dir/usr/share/$PRODUCT-release
                echo "$PRODUCT release $version_str" > $pkg_build_dir/etc/$PRODUCT-release
                ln -snf $PRODUCT-release $pkg_build_dir/etc/system-release
                sed -i "s|/openEuler-release|/$PRODUCT-release|g" $spec_file
                ;;
            "openEuler-logos")
                if [ -d "$LOGO_DIR" ]; then
                    cp -r $LOGO_DIR/* $pkg_build_dir
                fi
                mv $pkg_build_dir/usr/share/licenses/openEuler-logos $pkg_build_dir/usr/share/licenses/$PRODUCT-logos
                sed -i "s|openEuler|$PRODUCT|g" $pkg_build_dir/usr/share/licenses/$PRODUCT-logos/COPYING
                sed -i "s|/openEuler-logos|/$PRODUCT-logos|g" $spec_file
                ;;
            "openEuler-latest-release")
                mv $pkg_build_dir/etc/openEuler-latest $pkg_build_dir/etc/$PRODUCT-latest
                sed -i -r "s|^openeulerversion.*|${product}version=$OS_FULLNAME|g" $pkg_build_dir/etc/$PRODUCT-latest
                local dt=$(date +%F-%H-%M-%S)
                sed -i -r "s|(^compiletime=).*|\1$dt|g" $pkg_build_dir/etc/$PRODUCT-latest
                sed -i "s|openeulerversion|${product}version|g" $spec_file
                sed -i "s|/openEuler-latest|/$PRODUCT-latest|g" $spec_file
                sed -i "s|openEulerLinux.conf|${PRODUCT}Linux.conf|g" $spec_file
                ;;
            "anaconda-core")
                cp $ROOT_DIR/packages/anaconda/os.conf $pkg_build_dir/etc/anaconda/profile.d/$product.conf
                sed -i "s/OSNAME/$PRODUCT/g" $pkg_build_dir/etc/anaconda/profile.d/$product.conf
                rm -f $pkg_build_dir/etc/anaconda/profile.d/openeuler.conf
                sed -i "s|openeuler.conf|$product.conf|g" $spec_file
                ;;
            "python3-productmd")
                local treeinfo_file=$(find $pkg_build_dir -name treeinfo.py | head -n1)
                if [ -f $treeinfo_file ]; then
                    sed -i "s/openEuler/$PRODUCT/g" $treeinfo_file
                fi
                ;;
            "grub2-common")
                ## NOTE: replace openEuler to $PRODUCT, start failed, otherwise install failed

                ## install grub2-common failed (cpio: symlink failed) in UEFI, but ok in BIOS
                # ln -snf openEuler $pkg_build_dir/boot/efi/EFI/$PRODUCT
                # sed -i "/EFI\/openEuler\"/a%attr(0777, root, root) \"/boot/efi/EFI/$PRODUCT\"" $spec_file

                ## install failed (ln: command not found)
                # sed -i "/%changelog/i%post -p /bin/sh\nln -snf /boot/efi/EFI/openEuler /boot/efi/EFI/$PRODUCT\n\n" $spec_file

                ## change efi_dir in $product.conf for UEFI
                :
                ;;
        esac

        # rpmbuild
        rpmbuild -bb --buildroot=$pkg_build_dir $spec_file --nocheck --nosignature
        if [ $? -ne 0 ]; then
            log_error "rebuild rpm $pkg failed"
            return 1
        fi
    done

    # make iso dir
    log_info "make iso dir and create repo ..."
    clean_dir $ISO_DIR
    init_dir $ISO_DIR
    for i in $(ls $MOUNT_DIR); do
        if [[ "$i" = "repodata" ]]; then
            continue
        fi
        cp -ar $MOUNT_DIR/$i $ISO_DIR
    done
    mkdir -p $ISO_DIR/{Packages,repodata}
    find $rpm_dst_dir -name "*.rpm" -exec cp -a {} $ISO_DIR/Packages/ \;
    cp -a $MOUNT_DIR/repodata/normal.xml $ISO_DIR/repodata/normal.xml
    createrepo -g $ISO_DIR/repodata/normal.xml $ISO_DIR

    # replace label
    local cfg_files=(
        EFI/BOOT/grub.cfg
        isolinux/grub.conf
        isolinux/isolinux.cfg
    )
    for cfg in ${cfg_files[@]}; do
        local cfg_file=$ISO_DIR/$cfg
        sed -i "s|$oe_version|$os_version|g" $cfg_file
        sed -i "s|openEuler|$PRODUCT|g" $cfg_file
    done

    # custom install.img
    clean_dir $IMG_DIR
    init_dir $IMG_DIR

    local install_img_dir=$IMG_DIR/install_img
    local liveos_dir=$install_img_dir/squashfs-root/LiveOS
    local rootfs_img_dir=$IMG_DIR/rootfs_img
    local src_img=$MOUNT_DIR/images/install.img
    local dst_img=$ISO_DIR/images/install.img

    [[ -d $install_img_dir ]] && rm -rf $install_img_dir
    mkdir -p $install_img_dir $rootfs_img_dir

    cp -af $src_img $install_img_dir

    cd $install_img_dir
    unsquashfs install.img
    cd -

    cd $liveos_dir
    kpartx -av rootfs.img
    mount rootfs.img $rootfs_img_dir
    cd -

    cd $rootfs_img_dir

    cat << EOF > etc/os-release
NAME="$PRODUCT"
VERSION="$version_str"
ID="$PRODUCT"
VERSION_ID="$VERSION"
PRETTY_NAME="$PRODUCT $version_str"
ANSI_COLOR="0;31"
EOF

    mv usr/share/openEuler-release usr/share/$PRODUCT-release
    mv usr/share/licenses/openEuler-logos usr/share/licenses/$PRODUCT-logos
    if [ -d "$LOGO_DIR/usr/share/anaconda/pixmaps" ]; then
        cp -af $LOGO_DIR/usr/share/anaconda/pixmaps/*.png usr/share/anaconda/pixmaps
    fi
    cp -af $ROOT_DIR/packages/anaconda/os.conf etc/anaconda/profile.d/$product.conf
    sed -i "s/OSNAME/$PRODUCT/g" etc/anaconda/profile.d/$product.conf
    rm -f etc/anaconda/profile.d/openeuler.conf
    local treeinfo_file=$(find . -name treeinfo.py | head -n1)
    if [ -f $treeinfo_file ]; then
        sed -i "s/openEuler/$PRODUCT/g" $treeinfo_file
    fi
    # display in anaconda install UI
    sed -i "s|$oe_version|$os_version|g" .buildstamp
    sed -i "s|openEuler|$PRODUCT|g" .buildstamp

    # otherwise, umount target is busy
    cd -

    umount $rootfs_img_dir

    cd $install_img_dir
    [[ -f $dst_img ]] && rm -f $dst_img
    # select xz compression, default is gzip
    local bcj="x86"
    if [ "$ARCH" = "aarch64" ]; then
        bcj="arm"
    fi
    mksquashfs $install_img_dir/squashfs-root $dst_img -comp xz -Xbcj $bcj
    cd -

    # custom efiboot.img
    log_info "custom efiboot image ..."
    local efiboot_mount_dir=$IMG_DIR/efiboot_img
    local efiboot_img=$ISO_DIR/images/efiboot.img
    local efiboot_grub_cfg="$efiboot_mount_dir/EFI/BOOT/grub.cfg"
    if [ -d "$efiboot_mount_dir" ]; then
        umount $efiboot_mount_dir 2> /dev/null
    else
        mkdir -p $efiboot_mount_dir
    fi
    mount -o rw $efiboot_img $efiboot_mount_dir > /dev/null
    if [ $? -eq 0 ]; then
        sed -i "s|$oe_version|$os_version|g" $efiboot_grub_cfg
        sed -i "s|openEuler|$PRODUCT|g" $efiboot_grub_cfg
        umount $efiboot_mount_dir 2> /dev/null
    fi

    # remake iso
    log_info "remake iso image ..."
    make_iso $OS_FULLNAME-$ARCH $CUR_DIR/$ISO_OUTFILE $ISO_DIR
    if [ $? -ne 0 ]; then
        log_error "make iso failed"
        return 1
    fi

    # clean
    umount $MOUNT_DIR 2> /dev/null
    cd $CUR_DIR
    clean_dir $WORK_DIR

    return 0
}

main() {
    check_user

    parse_cmdline "$@"
    if [ $? -ne 0 ]; then
        log_error "parse params failed, try -h for help"
        return 1
    fi

    check_requires
    if [ $? -ne 0 ]; then
        log_error "check requires failed"
        return 1
    fi

    custom_iso
    if [ $? -ne 0 ]; then
        log_error "custom iso failed"
        return 1
    fi

    return 0
}

main "$@"
